import torch
from Net import Model
from transformers import BertTokenizer
from MyData import MyDataset
from torch.utils.data import DataLoader

DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 加载分词器
model_name = "./model/google-bert/bert-base-chinese/models--google-bert--bert-base-chinese/snapshots/c30a6ed22ab4564dc1e3b2ecbf6e766b0611a33f"
token = BertTokenizer.from_pretrained(model_name)

# 自定义数据编码处理函数
def collate_fn(data):
    sente = [i[0] for i in data]
    label = [i[1] for i in data]
    # 编码处理
    data = token.batch_encode_plus(
        batch_text_or_text_pairs=sente,
        truncation=True,
        padding='max_length',
        max_length=300,
        return_tensors='pt',
        return_length=True
    )
    input_ids = data['input_ids']
    attention_mask = data['attention_mask']
    token_type_ids = data['token_type_ids']
    labels = torch.LongTensor(label)

    return input_ids, attention_mask, token_type_ids, labels

# 创建数据集
train_dataset = MyDataset("train")
# 创建数据加载器
train_loader = DataLoader(
    dataset=train_dataset,
    batch_size=32,
    shuffle=True,
    drop_last=True,
    collate_fn=collate_fn  # 只传递函数名
)

if __name__ == "__main__":
    acc=0
    total=0
    model=Model().to(DEVICE)
    model.load_state_dict(torch.load("./params/1bert.pt"))
    model.eval() # 开启模型测试模式
    for i, (input_ids, attention_mask, token_type_ids, labels) in enumerate(train_loader):
        # 将数据放到DEVICE上
        input_ids, attention_mask, token_type_ids, labels = input_ids.to(DEVICE), attention_mask.to(
            DEVICE), token_type_ids.to(DEVICE), labels.to(DEVICE)
        # 执行前向计算
        out = model(input_ids, attention_mask, token_type_ids)

        out=out.argmax(dim=1)
        acc+=(out==labels).sum().item()
        total+=len(labels)
    print(acc/total)

# 0.8948958333333333